#!/usr/bin/python3
# -*- coding: utf-8 -*-
"""
@author:  runyuanye
@contact: runyuanye@163.com
"""

import argparse
import os
import sys
import random

sys.path.append('.')

"""
   对FaceID进行重排
"""


def main():
    parser = argparse.ArgumentParser(description="FaceID ReSort")
    parser.add_argument(
        "--faceid_file_list", default="/rootfs/media/kasim/Data1/data/ErisedPersonAge/SelectFace/RMRfaceid_file_list.txt", help="path to faceid file list", type=str
    )
    parser.add_argument(
        "--file_list", default="/rootfs/media/kasim/Data1/data/ErisedPersonAge/file_list.txt", help="path to file list", type=str
    )
    parser.add_argument(
        "--output_dir", default="/rootfs/media/kasim/Data1/data/ErisedPersonAge", help="path to output", type=str
    )
    args = parser.parse_args()

    random.seed(123456)

    output_dir = args.output_dir
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
        os.system('chmod a+wr {}'.format(output_dir))

    file_name_list = []
    with open(args.file_list, 'r') as file:
        for line in file.readlines():
            file_name = line.strip().split()[0]
            file_name_list.append(file_name)

    faceid = 1
    file_name_faceids = {}
    faceid_name_ids = {}
    with open(args.faceid_file_list, 'r') as file:
        for line in file.readlines():
            # 5269 5269 gd/zs/sqqjd/yxxhjt/150100414a54443452067fa6d4b06600/video/20190923_152901/00023_000.jpg 1
            lines = line.strip().split()
            faceid_name = '{}_{}'.format(lines[0], lines[1])
            file_name = lines[2]
            cur_faceid = faceid_name_ids.get(faceid_name, None)
            if cur_faceid is None:
                faceid_name_ids[faceid_name] = faceid
                cur_faceid = faceid
                faceid += 1
            file_name_faceids[file_name] = cur_faceid

    for file_name in file_name_list:
        if file_name not in file_name_faceids.keys():
            file_name_faceids[file_name] = faceid
            faceid += 1

    faceid_file_names = {}
    for file_name, faceid in file_name_faceids.items():
        file_names = faceid_file_names.get(faceid, None)
        if file_names is None:
            faceid_file_names[faceid] = [file_name]
        else:
            file_names.append(file_name)

    file_names = list(file_name_faceids.keys())
    file_names.sort()
    out_file_path = os.path.join(output_dir, 'file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        for file_name in file_names:
            faceid = file_name_faceids[file_name]
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    out_file_path = os.path.join(output_dir, 'typical_file_faceid_list.txt')
    with open(out_file_path, 'w') as file:
        typical_file_name_faceid = []
        for faceid, file_names in faceid_file_names.items():
            file_name = random.choice(file_names)
            typical_file_name_faceid.append([file_name, faceid])
        typical_file_name_faceid.sort(key=lambda x: x[0])
        for file_name, faceid in typical_file_name_faceid:
            out_info = '{} {}\n'.format(file_name, faceid)
            file.write(out_info)
    os.system('chmod a+wr {}'.format(out_file_path))

    print('FaceID Count:', len(faceid_file_names.keys()))
    print('finish!')


if __name__ == '__main__':
    main()
